import numpy as np
import torch
from torch import Tensor
from torch.utils.data import Dataset


class NNDataset(Dataset):
    def __getitem__(self, index):
        image_org, label_org = self._restore_data(index)
        image_aug, label_aug = self._augment(image_org, label_org)
        return self._build_input(image_aug), self._build_target(label_aug), label_aug

    def __len__(self):
        raise NotImplementedError

    def _restore_data(self, index):
        raise NotImplementedError

    def _build_input(self, image):
        return image.astype(np.float32).transpose([2, 0, 1])

    def _build_target(self, label):
        raise NotImplementedError

    def _augment(self, image, label):
        raise NotImplementedError

    @staticmethod
    def collate_fn(batch_data):
        """
        对批数据进行转置，从而使其转化成定长（与单个元素长度相同）输出,
        在此过程中，同类的数据将被聚集(collect)到一起.
        """
        batch = list(zip(*batch_data))
        to_collect, labels = batch[-2:]
        if isinstance(to_collect[0], (np.ndarray, Tensor)):
            collection = stack(to_collect)
        elif isinstance(to_collect[0], (tuple, list)):
            collection = [stack(target) for target in zip(*to_collect)]
        elif isinstance(to_collect[0], dict):
            collection = {
                sub_targets[0][0]: stack([sub_target for _, sub_target in sub_targets])
                for sub_targets in zip(*[target.items() for target in to_collect])
            }
        else:
            raise NotImplementedError
        if len(batch) == 3:
            return torch.from_numpy(np.stack(batch[0])), collection, labels
        return collection, labels


def stack(x):
    if isinstance(x[0], np.ndarray):
        return torch.from_numpy(np.stack(x))
    return torch.stack(x)
